#include <iostream>
#include <fstream>

#include "src/Trt.h"
#include "src/preprocess.h"
#include "src/postprocess.h"
#include "utils/viz.h"

#include <chrono>

using namespace std;


int main(int argc, char** argv)
{
    const std::string onnx_file ="../model/models.onnx";
    const char* img_file = "../img/street.jpg";
    const char* save_path = "../img/street_box.jpg";
    
    Trt trt;
    trt.create_engine(onnx_file);
    trt.print_info();
    void** buffer_io = trt.get_io();

    cv::Mat img = cv::imread(img_file);
    cv::Mat img_rgb;
    cv::cvtColor(img, img_rgb, cv::COLOR_BGR2RGB);
    preprocess(img_rgb, {640, 640}, (float*)(buffer_io[0]));

    trt.infer();

    std::vector<Bboxcv> result_obj_nms;
    Parameter para;
    para.head_shape = {20, 40, 80};
    para.input_shape = {640, 640};
    para.img_shape = {img.cols, img.rows};
    para.anchor = {142, 110, 192, 243, 459, 401, 36, 75, 76, 55, 72, 146, 12, 16, 19, 36, 40, 28};

    postprocess((float*)(buffer_io[3]), (float*)(buffer_io[2]), (float*)(buffer_io[1]), result_obj_nms, para);

    viz(img_file, save_path, result_obj_nms);
    return 0;
}